In `@src/megatron/bridge/models/qwen_omni/context_parallel_utils.py`:
- Around line 1-4: Add the missing Apache-2.0 / NVIDIA copyright header at the
top of the module context_parallel_utils.py (before the first import) so the
file follows project policy; update the header text consistent with other files
in the repo (include the Apache 2.0 license notice and NVIDIA copyright line),
and ensure there is a blank line after the header before the existing imports
(torch and megatron.core.parallel_state as mpu).
- Around line 193-201: The function split_data_cp_rank currently asserts on
val.shape before checking if val is None and declares cp_rank with an implicit
Optional; fix by making the None-check first (move the "if val is None: return
val" to the top of split_data_cp_rank before any access to val.shape) or, if val
should never be None, remove that guard and enforce non-None callers; also
update the signature from cp_rank: int = None to cp_rank: int | None and retain
the fallback cp_rank = mpu.get_context_parallel_rank() logic so behavior is
unchanged.

In `@src/megatron/bridge/models/qwen_omni/modelling_qwen3_omni.py`:
- Around line 41-42: The constructor accepts talker_transformer_config and
code2wav_transformer_config but never uses or stores them; either remove these
parameters from the __init__ signature or persist/forward them (e.g., assign to
self.talker_transformer_config and self.code2wav_transformer_config or pass them
into the appropriate submodule constructors) so callers' values take effect; if
they are placeholders, add a clear TODO comment in the class constructor noting
intended future use (reference the parameter names talker_transformer_config and
code2wav_transformer_config in the class constructor to locate the change).
- Around line 95-141: Change the positional forwarding in forward to explicit
keyword forwarding so arguments are passed by name to self.thinker (e.g.,
self.thinker(input_ids=input_ids, input_features=input_features,
position_ids=position_ids, ... , **kwargs)) to avoid brittle positional
coupling; also fix type annotations: change extra_block_kwargs: dict = None ->
dict | None = None, cp_img_num: list[int] = None -> list[int] | None = None, and
add missing optional types for parameters referenced in this method such as
input_features: torch.Tensor | None = None, inference_params: InferenceParams |
None = None, packed_seq_params: PackedSeqParams | None = None,
feature_attention_mask: torch.Tensor | None = None, audio_feature_lengths:
torch.Tensor | None = None, use_audio_in_video: bool | None = None, and
video_second_per_grid: float | None = None so static analysis (RUF013) is
satisfied while keeping parameter names matching those passed to self.thinker.

In `@src/megatron/bridge/models/qwen_omni/qwen3_omni_bridge.py`:
- Around line 15-35: Remove the unused imports at the top of
qwen3_omni_bridge.py: drop logging, Dict/Mapping/Union from typing, torch.nn,
megatron.core.parallel_state, WeightConversionTask, and
extract_expert_number_from_param; keep only the actual used symbols such as
Qwen3OmniMoeForConditionalGeneration, MegatronMappingRegistry,
Qwen3OmniMoeModel, MegatronModelBridge, the param mapping classes (AutoMapping,
ConcatenatedQKVMapping, GatedMLPMapping, QKVMapping, ReplicatedMapping),
PreTrainedVLM, and Qwen3OmniMoeModelProvider so imports reflect real usage in
this module.

In `@src/megatron/bridge/models/qwen_omni/qwen3_omni_provider.py`:
- Around line 144-146: The docstring for the provide method incorrectly says
"Qwen3VL MoE"; update the string to "Qwen3 Omni MoE" to accurately describe the
model. Locate the provide method in qwen3_omni_provider.py (function provide)
and replace the phrase "Qwen3VL MoE" with "Qwen3 Omni MoE" in its docstring,
ensuring the rest of the description remains unchanged.
- Line 24: Replace the old typing imports and annotations: remove "from typing
import List, Optional" and update any annotations using List[...] to use
built-in list[...] and Optional[X] to use X | None (e.g., change List[str] ->
list[str], Optional[int] -> int | None). Apply these edits to the import line
and all type hints in the Qwen3OmniProvider implementation (fields and function
signatures referenced around the earlier import and at the sections identified
near lines 41–44, 84, and 108) so the module uses Python 3.10+ built-in generics
and the T | None nullable syntax.

In `@src/megatron/bridge/models/qwen_omni/thinker_model.py`:
- Around line 7-11: The file header in thinker_model.py has an incomplete Apache
2.0 license block; update the top-of-file license to match the complete header
used elsewhere (e.g., utils.py) by adding the missing sentences "distributed
under the License is distributed on an 'AS IS' BASIS, WITHOUT WARRANTIES OR
CONDITIONS OF ANY KIND, either express or implied." so the header includes the
full Apache clause and mirrors the other files' license text.
- Around line 323-342: The audio embeddings variable audio_embeds is never set
after computing audio_features, so audio modality is dropped; assign
audio_embeds from the computed audio_features (e.g., set audio_embeds =
audio_features or convert audio_features into the same shape expected by
combined_embeddings) after calling get_audio_features in the block where
input_features is not None, ensure audio_mask is computed using input_ids ==
self.audio_token_id, and then the existing conditional that writes
combined_embeddings[audio_mask] = audio_embeds will correctly insert the audio
embeddings; reference functions/variables: get_audio_features, audio_features,
audio_embeds, audio_mask, combined_embeddings, audio_feature_lengths.
- Line 13: Remove the Optional import and switch all nullable type annotations
to PEP 604 unions (X | None); specifically update the function signatures that
currently use implicit Optional defaults (e.g., parameters declared as "dict =
None" and "list[int] = None") to use explicit union annotations like "dict |
None" and "list[int] | None", and update get_audio_features signature to use the
same pattern; ensure you remove the now-unused Optional import from the top of
thinker_model.py and update any affected type hints accordingly.
- Around line 205-211: The code can raise AttributeError when both masks are
None: in get_audio_features the branch sets audio_feature_lengths = None and
then computes feature_lens = audio_feature_lengths if ... else
feature_attention_mask.sum(-1), which calls .sum on None; fix by making the else
path compute feature_lens from input_features instead of feature_attention_mask
(e.g. set feature_lens = torch.full((input_features.size(0),),
input_features.size(-1), dtype=torch.long, device=input_features.device) or
similar) and only use feature_attention_mask.sum(-1) when feature_attention_mask
is not None; update the logic around audio_feature_lengths/feature_lens to use
input_features, feature_attention_mask, and audio_feature_lengths safely.
- Around line 294-298: The fallback zero-tensor uses a hardcoded dtype
(torch.bfloat16); change it to use the language model's actual dtype so it
matches runtime precision. In the method creating vision_embeds (look for the
vision_embeds = torch.zeros(...) fallback in thinker_model.py) obtain the model
dtype via a safe lookup such as: model_dtype = getattr(self.language_model,
"dtype", None) or getattr(self.language_model.config, "torch_dtype", None) or
fallback to next(self.language_model.parameters()).dtype, then create the zeros
with dtype=model_dtype; apply the same change to the other fallback zero-tensor
occurrence so both use the model's dtype rather than torch.bfloat16.
- Around line 141-145: The assert currently passes a tuple as the message which
prints oddly; update the assert in thinker_model.py that references
vision_transformer_config.vision_config.deepstack_visual_indexes and
self.language_model.decoder.layers so the second argument is a single formatted
string (e.g., an f-string) combining the descriptive texts and values instead of
a tuple; keep the same condition and include both counts in the single message
for clear output.

In `@src/megatron/bridge/models/qwen_omni/transformer_config.py`:
- Around line 15-16: The file imports and type annotations use typing.List and
typing.Optional; update the import and annotations in transformer_config.py to
use built-in generics and PEP 604 union syntax: replace any "List[T]" with
"list[T]" and "Optional[T]" with "T | None" (remove List and Optional from the
typing import), and update all corresponding fields/usages in the dataclass
(e.g., the TransformerConfig dataclass and the other three annotated attributes
referenced in the review) to use the new forms.
- Around line 18-23: Replace the direct import of TransformerConfig from
megatron.core.transformer.transformer_config with the bridge wrapper import used
across providers: import TransformerConfig from
megatron.bridge.models.transformer_config so Qwen3OmniTransformerConfig inherits
the bridge's wrapper class (update the import statement that currently
references TransformerConfig and keep Qwen3OmniTransformerConfig as-is).

In `@src/megatron/bridge/models/qwen_omni/utils.py`:
- Line 16: Replace the use of typing.Optional with Python 3.10 union syntax:
remove the Optional import and update any function/type hints that use
Optional[...] to use X | None instead; specifically change the get_rope_index
signature and any other occurrences in the same file (around the block
referenced at lines 60-66) to use the `Type | None` form (e.g., replace
Optional[int] with int | None) and ensure imports reflect removal of Optional.
- Around line 86-88: The loop reuses the outer-scope name input_ids (for i,
input_ids in enumerate(total_input_ids)), which gets masked and then later used
incorrectly; rename the loop variable to _input_ids (for i, _input_ids in
enumerate(total_input_ids)), apply the attention mask to a local masked_ids =
_input_ids[_mask == 1] (or similar) and replace all downstream references inside
the loop from input_ids to the new name (e.g., masked_ids or _input_ids) so you
no longer shadow the outer variable; also change the places that rely on the
original tensor/device/length (notably the code that uses len(input_ids) and
input_ids.device around the delta calculation) to use total_input_ids (or
total_input_ids.device / len(total_input_ids)) instead.

In `@src/megatron/bridge/recipes/qwen_vl/qwen3_vl.py`:
- Line 251: The hf_path in the config is a local relative path (hf_path) instead
of the HuggingFace hub identifier expected by AutoBridge.from_hf_pretrained;
replace "../hf-hub/Qwen/Qwen3-Omni-30B-A3B-Instruct" with the official HF repo
id string (e.g., "Qwen/Qwen3-Omni-30B-A3B-Instruct" or the correct hub name used
elsewhere) so AutoBridge.from_hf_pretrained can resolve the model from the hub;
update the hf_path value in the qwen3_vl.py config to use the hub identifier.